import pathlib

import matplotlib.pyplot
import torch

from checkpoint import Checkpoint
from checkpoint_manager import CheckpointManager
from scaler import Scaler
from lotka_volterra_dataset import LotkaVolterraDataset
from lvminn import Lvminn


def main():
    dataset = LotkaVolterraDataset.load(pathlib.Path("./checkpoints/data.pt"))
    t_min = float(dataset.t_min())
    t_max = float(dataset.t_max())

    model = Lvminn(Scaler(t_min, t_max))

    checkpoints = CheckpointManager(pathlib.Path("./checkpoints")).existed()
    checkpoints = sorted(checkpoints.items(), key=lambda x: x[0])
    checkpoints_step = len(checkpoints) // 100
    if checkpoints_step > 0:
        checkpoints = checkpoints[checkpoints_step - 1::checkpoints_step]

    matplotlib.pyplot.figure(figsize=(8, 5), dpi=110)

    scatter_a1 = None
    scatter_b1 = None
    scatter_a2 = None
    scatter_b2 = None
    for epoch, checkpoint_path in checkpoints:
        checkpoint = Checkpoint.load(checkpoint_path)
        model.load_state_dict(checkpoint.model_state())

        matplotlib.pyplot.subplot(221)
        scatter_a1 = matplotlib.pyplot.scatter(
            epoch, torch.detach(model.a1),
            c="tab:orange", s=4)

        matplotlib.pyplot.subplot(222)
        scatter_b1 = matplotlib.pyplot.scatter(
            epoch, torch.detach(model.b1),
            c="tab:orange", s=4)

        matplotlib.pyplot.subplot(223)
        scatter_a2 = matplotlib.pyplot.scatter(
            epoch, torch.detach(model.a2),
            c="tab:orange", s=4)

        matplotlib.pyplot.subplot(224)
        scatter_b2 = matplotlib.pyplot.scatter(
            epoch, torch.detach(model.b2),
            c="tab:orange", s=4)

    matplotlib.pyplot.subplot(221)
    line_a1 = matplotlib.pyplot.axhline(dataset.problem().a1(), c="tab:green")
    matplotlib.pyplot.legend(
        [line_a1, scatter_a1],
        [f"Expected  = {dataset.problem().a1():.2e}", f"Estimated = {float(model.a1):.2e}"])
    matplotlib.pyplot.ylabel(r"$a_1$")
    matplotlib.pyplot.xticks([100000, 200000])

    matplotlib.pyplot.subplot(222)
    line_b1 = matplotlib.pyplot.axhline(dataset.problem().b1(), c="tab:green")
    matplotlib.pyplot.legend(
        [line_b1, scatter_b1],
        [f"Expected  = {dataset.problem().b1():.2e}", f"Estimated = {float(model.b1):.2e}"])
    matplotlib.pyplot.ylabel(r"$b_1$")
    matplotlib.pyplot.xticks([100000, 200000])

    matplotlib.pyplot.subplot(223)
    line_a2 = matplotlib.pyplot.axhline(dataset.problem().a2(), c="tab:green")
    matplotlib.pyplot.legend(
        [line_a2, scatter_a2],
        [f"Expected  = {dataset.problem().a2():.2e}", f"Estimated = {float(model.a2):.2e}"])
    matplotlib.pyplot.ylabel(r"$a_2$")
    matplotlib.pyplot.xlabel("Epoch")
    matplotlib.pyplot.xticks([100000, 200000])

    matplotlib.pyplot.subplot(224)
    line_b2 = matplotlib.pyplot.axhline(dataset.problem().b2(), c="tab:green")
    matplotlib.pyplot.legend(
        [line_b2, scatter_b2],
        [f"Expected  = {dataset.problem().b2():.2e}", f"Estimated = {float(model.b2):.2e}"])
    matplotlib.pyplot.ylabel(r"$b_2$")
    matplotlib.pyplot.xlabel("Epoch")
    matplotlib.pyplot.xticks([100000, 200000])

    matplotlib.pyplot.subplots_adjust(wspace=0.35)
    matplotlib.pyplot.show()


if __name__ == "__main__":
    main()
